#include <arch/lapic.h>
#include <arch/io.h>
#include <arch/processor.h>

#include <drv/acpi.h>
#include <drv/serial.h>
#include <drv/hpet.h>

#include <cpuid.h>

#define MADT_LAPIC (0)
#define MADT_IOAPIC (1)
#define MADT_IOAPIC_INTSRC_OVERRIDE (2)
#define MADT_IOAPIC_NMI_SRC (3)
#define MADT_LAPIC_NMI_SRC (4)
#define MADT_LAPIC_ADDR_OVERRIDE (5)
#define MADT_LX2APIC (6)

uintptr_t lapic_base = 0;
uintptr_t ioapic_base = 0;

typedef struct __packed {
	uint8_t type;
	uint8_t length;
} madt_entry_t;

typedef struct __packed {
	acpi_table_header_t header;
	uint32_t lapic_base;
	uint32_t flags;
	madt_entry_t entries[0];
} acpi_madt_t;

typedef struct __packed {
	madt_entry_t header;
	uint8_t acpi_processor_id;
	uint8_t apic_id;
	uint32_t flags;
} madt_lapic_t;

typedef struct __packed {
	madt_entry_t header;
	uint8_t ioapic_id;
	uint8_t reserved;
	uint32_t ioapic_base;
	uint32_t global_system_interrupt_base;
} madt_ioapic_t;

typedef struct __packed {
	madt_entry_t header;
	uint16_t reserved;
	uint32_t x2apic_id;
	uint32_t flags;
	uint32_t acpi_processor_id;
} madt_lx2apic_t;

void lapic_init() {
	acpi_madt_t *madt = (acpi_madt_t *)acpi_get_table("APIC");

	lapic_base = madt->lapic_base;

	uint32_t length = madt->header.length - sizeof(acpi_madt_t);
	for(madt_entry_t *mentry = madt->entries; length; length -= mentry->length, mentry = (madt_entry_t *)((uintptr_t)mentry + mentry->length)) {
		switch(mentry->type) {
			case MADT_IOAPIC:
				ioapic_base = ((madt_ioapic_t *)mentry)->ioapic_base;
				break;
			case MADT_LAPIC:
				madt_lapic_t *lapic = (madt_lapic_t *)mentry;
				processor_register(lapic->acpi_processor_id, lapic->apic_id, lapic->flags);
				break;
			case MADT_LX2APIC:
				madt_lx2apic_t *lx2apic = (madt_lx2apic_t *)mentry;
				processor_register(lx2apic->acpi_processor_id, lx2apic->x2apic_id, lx2apic->flags);
				break;
		}
	}

	io_write8(0x21, 0xff);
	io_write8(0xa1, 0xff);

	for(uint64_t i = 0; i < 6; ++i) {
		lapic_write_lvt(i, 0x10000);
	}

	uint64_t apic_msr = io_rdmsr(0x1b);

	uint32_t ecx, tmp;
	__cpuid(1, tmp, tmp, ecx, tmp);
	if(ecx & BIT(21)) {
		apic_msr |= BIT(10);	// 启用x2apic
	}

	io_wrmsr(0x1b, apic_msr | BIT(11));
	apic_msr = io_rdmsr(0x1b);

	if(!(apic_msr & BIT(11))) {
		__asm__ volatile("hlt");
	}

	if(apic_msr & BIT(10)) {
		lapic_base = 0;
	}

	lapic_write32(0xf0, 0x1ff);

	processor_get_current()->attribute |= PROCESSOR_BSP;
}

uint32_t lapic_ap_init() {
	uint64_t apic_msr = io_rdmsr(0x1b);

	uint32_t ecx, tmp;
	__cpuid(1, tmp, tmp, ecx, tmp);
	if(ecx & BIT(21)) {
		apic_msr |= BIT(10);	// 启用x2apic
	}

	io_wrmsr(0x1b, apic_msr | BIT(11));
	apic_msr = io_rdmsr(0x1b);

	if(!(apic_msr & BIT(11))) {
		__asm__ volatile("hlt");
	}

	if(apic_msr & BIT(10)) {
		lapic_base = 0;
	}

	for(uint64_t i = 0; i < 6; ++i) {
		lapic_write_lvt(i, 0x10000);
	}

	lapic_write32(0xf0, 0x1ff);

	return lapic_get_self_id();
}

void lapic_write32(uint32_t reg, uint32_t val) {
	if(lapic_base) {
		mmio_write32(lapic_base, reg, val);
	} else {
		io_wrmsr(0x800 + (reg >> 4), val);
	}
}

void lapic_write64(uint32_t reg, uint64_t val) {
	if(lapic_base) {
		mmio_write32(lapic_base, reg + 0x10, val >> 32);
		mmio_write32(lapic_base, reg, val);
	} else {
		io_wrmsr(0x800 + (reg >> 4), val);
	}
}

uint32_t lapic_read32(uint32_t reg) {
	if(lapic_base) {
		return mmio_read32(lapic_base, reg);
	} else {
		return io_rdmsr(0x800 + (reg >> 4));
	}
}

uint64_t lapic_read64(uint32_t reg) {
	if(lapic_base) {
		return ((uint64_t)mmio_read32(lapic_base, reg + 0x10) << 32) | mmio_read32(lapic_base, reg);
	} else {
		return io_rdmsr(0x800 + (reg >> 4));
	}
}

uint32_t lapic_read_lvt(uint32_t index) {
	return lapic_read32(0x320 + (index << 4));
}

void lapic_write_lvt(uint32_t index, uint32_t val) {
	lapic_write32(0x320 + (index << 4), val);
}

void lapic_mask_lvt(uint32_t index, bool enable) {
	uint32_t lvt = lapic_read_lvt(index);

	lapic_write_lvt(index, enable ? lvt & ~BIT(16) : lvt | BIT(16));
}

void lapic_send_ipi(uint32_t apic_id, uint32_t ipi_type, uint8_t vector) {
	if(lapic_base) {
		apic_id <<= 24;
	}

	lapic_write64(0x300, ((uint64_t)apic_id << 32) | ((ipi_type << 8) & 0xffffff00) | vector);

	if(lapic_base) {
		while(lapic_read32(0x300) & BIT(12)) {
			__asm__ volatile("pause");
		}
	}
}

uint32_t lapic_get_self_id() {
	uint32_t id = lapic_read32(0x20);

	return lapic_base ? id >> 24 : id;
}

void lapic_start_timer() {
	__asm__ volatile("cli");

	lapic_write32(0x3e0, 0b1011);	 // 除数为1
	lapic_write32(0x380, ~0u);		 // 清零计时器

	hpet_sleep(APIC_TIMER_INTERVAL);

	uint32_t tick = 0xffffffff - lapic_read32(0x390);

	lapic_write_lvt(LVT_IDX_TIMER, 0x20 | BIT(17));	 // 时钟中断为irq 0x20

	lapic_write32(0x3e0, 0b1011);	 // 再次设置除数为1

	__asm__ volatile("sti");

	lapic_write32(0x380, tick);	 // 开始计时
}